import csv
import math
import warnings
import pandas as pd
from sklearn.model_selection import GridSearchCV
from xgboost import XGBRegressor

warnings.filterwarnings("ignore", category=UserWarning)
dic = {}
ref = 0


def get_data(f):
    global dic, ref
    d = pd.read_csv(f)
    tmphead = list(d.head())
    tmplist = list(d.values)
    rlist = []
    for i in tmplist:
        tmp = []
        for t in i:
            try:
                if (math.isnan(float(t))):
                    try:
                        tmp.append(dic['NA'])
                    except:
                        dic['NA'] = ref
                        ref += 1
                        tmp.append(dic['NA'])
                else:
                    tmp.append(float(t))
            except:
                try:
                    tmp.append(dic[t])
                except:
                    dic[t] = ref
                    ref += 1
                    tmp.append(dic[t])
        rlist.append(tmp)
    return pd.DataFrame(rlist, columns=tmphead)


def get_target(f):
    global dic
    d = pd.read_csv(f)
    tmphead = list(d.head())
    tmplist = list(d.values)
    rlist = []
    for i in tmplist:
        tmp = []
        for t in i:
            try:
                if (math.isnan(float(t))):
                    tmp.append(dic['NA'])
                else:
                    tmp.append(float(t))
            except:
                tmp.append(dic[t])
        rlist.append(tmp)
    return pd.DataFrame(rlist, columns=tmphead)


def writ_data(idi, data, file):
    with open(file, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['ID', 'SalePrice'])
        for i in range(len(data)):
            writer.writerow([int(idi[i]), data[i]])


def get_result(my_model, target):
    return my_model.predict(target)


def select_model(train, label):
    params = {
        'booster': ['gbtree', 'gblinear', 'dart'],
        'learning_rate': [0.01, 0.1, 0.2],
        'n_estimators': [50, 100, 150]

    }
    model = XGBRegressor()
    my_model = GridSearchCV(model, params, cv=5, error_score=1, refit=True, n_jobs=-1)
    my_model.fit(train, label)
    return my_model, my_model.best_score_


if __name__ == "__main__":
    d = get_data('./data/train.csv')
    corrmat = d.corr()
    rela = list(corrmat['SalePrice'].abs().sort_values().index)[:-1]
    features = 68
    select_feat = rela[-features:]

    train, label = d.drop(['SalePrice'], axis=1, inplace=False), d['SalePrice']
    train = train[select_feat]

    d = get_data('./data/test.csv')
    idi = d['Id']
    target = d.drop(['Id'], axis=1, inplace=False)
    target = target[select_feat]
    my_model, score = select_model(train, label)
    print(my_model)
    print(score)

    result = get_result(my_model, target)
    writ_data(idi, result, './data/result.csv')
